"""
This script serves as the main entry point for running evaluation experiments
using Hydra for configuration management. It supports single runs and
multirun sweeps over various parameters, including 'inference_horizon'.

To run a single experiment:
python scripts/run_eval.py ckpt_path=/path/to/your/checkpoint.msgpack inference_horizon=7

To run a multirun sweep over 'inference_horizon' (from 1 to 8):
(Ensure configs/sweeper/basic.yaml is configured with inference_horizon: 1,2,3,4,5,6,7,8)
python scripts/run_eval.py --multirun ckpt_path=/path/to/your/checkpoint.msgpack

To run a multirun sweep over 'inference_horizon' and a specific inference strategy (e.g., 'replan'):
(First, create configs/inference/replan.yaml and add the 'inference_replan' method to agent.py)
python scripts/run_eval.py --multirun ckpt_path=/path/to/your/checkpoint.msgpack inference_horizon=1,2,3,4,5,6,7,8 inference=replan

Results for each run will be saved in the 'outputs/' directory, organized by timestamp and run parameters.
"""

import logging

import flax.serialization
import hydra
import jax
import ogbench
from omegaconf import DictConfig, OmegaConf

from ne_flow import agent as agent_lib
from ne_flow import evaling

log = logging.getLogger(__name__)


@hydra.main(config_path="../configs", config_name="config", version_base=None)
def main(cfg: DictConfig) -> float:
    """
    Main evaluation function driven by Hydra configuration.
    """
    log.info("--- Configuration ---")
    log.info(OmegaConf.to_yaml(cfg))
    log.info("---------------------")

    # --- 1. Initialize Environment and Agent ---
    env, _, _ = ogbench.make_env_and_datasets(
        cfg.dataset_name,
        dataset_dir="data_static/",
        compact_dataset=True,
    )

    observation_shape = env.observation_space.shape
    action_shape = env.action_space.shape

    rng = jax.random.PRNGKey(cfg.seed)
    rng, agent_rng = jax.random.split(rng)

    # For now, we assume GCRLAgent0 is the one we want to test.
    # This could also be made configurable in the future.
    agent_class = agent_lib.GCRLAgent0

    # We need to construct a dictionary for the agent's config from the flat Hydra cfg
    # and dynamically build the encoder_config.
    agent_config_params = OmegaConf.to_container(cfg, resolve=True)

    # Remove keys not meant for the agent's internal config
    # These are top-level script parameters or handled separately.
    agent_config_params.pop('ckpt_path', None)
    agent_config_params.pop('num_episodes', None)
    agent_config_params.pop('inference', None) # Inference config is merged below
    agent_config_params.pop('encoder', None) # Encoder config is handled separately
    agent_config_params.pop('dataset_name', None)
    agent_config_params.pop('seed', None)

    # Add inference-specific params to the main agent config for self-contained agent
    agent_config_params.update(OmegaConf.to_container(cfg.inference, resolve=True))

    # --- Dynamically build encoder_config based on cfg.encoder ---
    input_dim = observation_shape[0]
    encoder_name = cfg.encoder.name

    if encoder_name == "identity":
        output_dim = input_dim
    else:
        output_dim = cfg.encoder.latent_dim

    encoder_config = [
        ("proprio", encoder_name, input_dim, output_dim),
    ]
    size_latent_obs = output_dim # The actual latent size

    # Also update cfg.encoder.latent_dim to reflect the actual output_dim for identity
    # This ensures consistency if other parts of the config rely on it.
    if encoder_name == "identity":
        cfg.encoder.latent_dim = output_dim

    agent_kwargs = {
        "config": agent_config_params,
        "encoder_config": encoder_config, # Use dynamically built config
        "size_latent_obs": size_latent_obs, # Pass the actual latent size
        "size_act": action_shape[0],
        "rng": agent_rng,
    }

    agent = agent_class(**agent_kwargs)

    # --- 2. Load Checkpoint ---
    log.info(f"Loading checkpoint from {cfg.ckpt_path}...")
    with open(cfg.ckpt_path, "rb") as f:
        loaded_bytes = f.read()

    loaded_params = flax.serialization.from_bytes(
        agent.combined_state.ema_params, loaded_bytes
    )
    agent.combined_state = agent.combined_state.replace(ema_params=loaded_params)

    # --- 3. Resolve and Bind Inference Strategy ---
    log.info(f"Resolving inference strategy: '{cfg.inference.name}'")
    method_name = cfg.inference.name
    try:
        inference_fn = getattr(agent, method_name)
    except AttributeError:
        log.error(
            f"Error: Agent {agent_class.__name__} has no method named '{method_name}'"
        )
        raise

    # All configuration is now inside the agent, so we don't need to bind any
    # parameters here. The agent's method is called directly.
    bound_inference_fn = inference_fn
    log.info(f"Using inference method: '{method_name}'")

    # --- 4. Run Evaluation ---
    rng, eval_rng = jax.random.split(rng)

    # We use do_eval_stats to get performance metrics
    overall_success, success_length = evaling.do_eval_stats(
        env,
        agent,
        num_eval_episode=cfg.num_episodes,
        inference_horizon=cfg.inference_horizon,
        rng=eval_rng,
        inference_fn=bound_inference_fn,  # Pass the resolved function
    )

    log.info(f"--- Results for inference_horizon={cfg.inference_horizon} ---")
    log.info(f"Overall Success Rate: {overall_success:.4f}")
    log.info(f"Average Length of Successful Episodes: {success_length:.2f}")
    log.info("-----------------------------------------------------")

    # Hydra can automatically log the return value of the main function
    return overall_success


if __name__ == "__main__":
    main()
